import json
import os
from pycocotools.coco import COCO
import matplotlib.pyplot as plt
from PIL import Image

json_path = "dataset/annotations/captions_val2014.json"
img_path = "dataset/val2014"

coco = COCO(annotation_file=json_path)

ids = list(sorted(coco.imgs.keys()))

print("number of images: {}".format(len(ids)))

ind = 1
for img_id in ids[:3000]:
    ann_ids = coco.getAnnIds(imgIds=img_id)

    targets = coco.loadAnns(ann_ids)
    prompt = []

    with open("dataset/test_coco_sr3_caption/COCO_val2014_"+str(img_id).zfill(12)+".txt", 'w') as f:
    # with open("dataset/test_coco_sr3_caption/0_"+str(ind)+"_hr.txt", 'w') as f:
        # for i in range(5):
        #     f.writelines(targets[i]["caption"]+'\n')
        f.writelines(targets[4]["caption"]+'\n')

    print(targets[4]["caption"])

    ind = ind + 1
    # path = coco.loadImgs(img_id)[0]['file_name']
    # print(path)
    # img = Image.open(os.path.join(img_path, path)).convert('RGB')

